# utils/training_monitor.py
from typing import Dict
import matplotlib.pyplot as plt
import numpy as np
from datetime import datetime
import json
import os

class TrainingMonitor:
    """训练监控器 - 跟踪和可视化训练进度"""
    
    def __init__(self):
        self.metrics_history = {
            'total_loss': [],
            'policy_losses': [],
            'value_losses': [],
            'performance': [],
            'rewards': [],
            'survival_rates': []
        }
        self.factor_metrics = {}
    
    def update_metrics(self, cycle_metrics: Dict):
        """更新指标历史"""
        # 总体指标
        self.metrics_history['total_loss'].append(cycle_metrics.get('total_loss', 0))
        self.metrics_history['performance'].append(cycle_metrics.get('overall_performance', 0))
        
        # 回合指标
        if 'episode_results' in cycle_metrics:
            episodes = cycle_metrics['episode_results']
            self.metrics_history['rewards'].append(np.mean([ep['total_reward'] for ep in episodes]))
            self.metrics_history['survival_rates'].append(
                np.mean([ep['surviving_factors'] for ep in episodes]) / 12.0
            )
        
        # 因子特定指标
        if 'training_metrics' in cycle_metrics:
            training_metrics = cycle_metrics['training_metrics']
            if 'policy_losses' in training_metrics:
                self._update_factor_metrics(training_metrics['policy_losses'], 'policy_loss')
            if 'value_losses' in training_metrics:
                self._update_factor_metrics(training_metrics['value_losses'], 'value_loss')
    
    def _update_factor_metrics(self, metrics: Dict, metric_type: str):
        """更新因子特定指标"""
        for factor_type, value in metrics.items():
            factor_name = factor_type.name
            if factor_name not in self.factor_metrics:
                self.factor_metrics[factor_name] = {
                    'policy_loss': [],
                    'value_loss': []
                }
            self.factor_metrics[factor_name][metric_type].append(value)
    
    def generate_training_report(self, cycle: int) -> str:
        """生成训练报告"""
        report = f"""
训练报告 - 第 {cycle} 轮回
生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

总体指标:
- 平均性能: {np.mean(self.metrics_history['performance'][-10:]):.3f}
- 平均奖励: {np.mean(self.metrics_history['rewards'][-10:]):.3f}
- 平均生存率: {np.mean(self.metrics_history['survival_rates'][-10:]):.3f}
- 当前损失: {self.metrics_history['total_loss'][-1]:.6f}

因子表现:
"""
        
        # 添加因子特定指标
        for factor_name, metrics in self.factor_metrics.items():
            if metrics['policy_loss']:
                avg_policy_loss = np.mean(metrics['policy_loss'][-5:])
                report += f"- {factor_name}: 策略损失 {avg_policy_loss:.6f}\n"
        
        return report
    
    def plot_training_progress(self, save_path: str = "training_plots"):
        """绘制训练进度图"""
        os.makedirs(save_path, exist_ok=True)
        
        # 创建图表
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))
        
        # 性能图表
        if self.metrics_history['performance']:
            ax1.plot(self.metrics_history['performance'])
            ax1.set_title('训练性能')
            ax1.set_xlabel('轮回')
            ax1.set_ylabel('性能')
            ax1.grid(True)
        
        # 损失图表
        if self.metrics_history['total_loss']:
            ax2.plot(self.metrics_history['total_loss'])
            ax2.set_title('总损失')
            ax2.set_xlabel('轮回') 
            ax2.set_ylabel('损失')
            ax2.grid(True)
        
        # 奖励图表
        if self.metrics_history['rewards']:
            ax3.plot(self.metrics_history['rewards'])
            ax3.set_title('平均奖励')
            ax3.set_xlabel('轮回')
            ax3.set_ylabel('奖励')
            ax3.grid(True)
        
        # 生存率图表
        if self.metrics_history['survival_rates']:
            ax4.plot(self.metrics_history['survival_rates'])
            ax4.set_title('平均生存率')
            ax4.set_xlabel('轮回')
            ax4.set_ylabel('生存率')
            ax4.grid(True)
        
        plt.tight_layout()
        plt.savefig(f"{save_path}/training_progress.png", dpi=300, bbox_inches='tight')
        plt.close()
        
        # 因子特定图表
        if self.factor_metrics:
            self._plot_factor_metrics(save_path)
    
    def _plot_factor_metrics(self, save_path: str):
        """绘制因子特定指标"""
        fig, axes = plt.subplots(2, 1, figsize=(12, 8))
        
        # 策略损失
        for factor_name, metrics in self.factor_metrics.items():
            if metrics['policy_loss']:
                axes[0].plot(metrics['policy_loss'], label=factor_name)
        axes[0].set_title('因子策略损失')
        axes[0].set_xlabel('更新步骤')
        axes[0].set_ylabel('策略损失')
        axes[0].legend()
        axes[0].grid(True)
        
        # 价值损失
        for factor_name, metrics in self.factor_metrics.items():
            if metrics['value_loss']:
                axes[1].plot(metrics['value_loss'], label=factor_name)
        axes[1].set_title('因子价值损失')
        axes[1].set_xlabel('更新步骤')
        axes[1].set_ylabel('价值损失')
        axes[1].legend()
        axes[1].grid(True)
        
        plt.tight_layout()
        plt.savefig(f"{save_path}/factor_metrics.png", dpi=300, bbox_inches='tight')
        plt.close()
    
    def save_metrics(self, filepath: str = "training_metrics.json"):
        """保存指标数据"""
        data = {
            'metrics_history': self.metrics_history,
            'factor_metrics': self.factor_metrics,
            'last_updated': datetime.now().isoformat()
        }
        
        with open(filepath, 'w') as f:
            json.dump(data, f, indent=2)
